#include <Uefi.h>
#include <Library/BaseLib.h>
#include <Library/PrintLib.h>
#include <Library/MemoryAllocationLib.h>

EFI_STATUS
EFIAPI
UefiMain(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
{
	BASE_LIBRARY_JUMP_BUFFER *JmpBuffer    = NULL;
	BASE_LIBRARY_JUMP_BUFFER *NewJmpBuffer = NULL;
	UINTN SetJmpFlag                       = 0;

	JmpBuffer = AllocatePool(sizeof(BASE_LIBRARY_JUMP_BUFFER) + BASE_LIBRARY_JUMP_BUFFER_ALIGNMENT);
	if (JmpBuffer == NULL) {
		Print(L"AllocatePool failed!\n");
		return EFI_OUT_OF_RESOURCES;
	}
	NewJmpBuffer = ALIGN_POINTER(JmpBuffer, BASE_LIBRARY_JUMP_BUFFER_ALIGNMENT);
	SetJmpFlag   = SetJump(NewJmpBuffer);

AfterJump:
	if (SetJmpFlag == 0) {
		Print(L"Returned from LongJump(), SetJmpFlag = %d\n", SetJmpFlag);
		LongJump(NewJmpBuffer, 1); // 跳转到 SetJump()位置，并返回 1
	} else if (SetJmpFlag == 1) {
		Print(L"Returned from LongJump(), SetJmpFlag = %d\n", SetJmpFlag);		
		LongJump(NewJmpBuffer, 2); // 跳转到 SetJump()位置，并返回 2
	}

	Print(L"Returned from LongJump(), SetJmpFlag = %d\n", SetJmpFlag);
	FreePool(JmpBuffer);
	return EFI_SUCCESS;
}